#!/usr/bin/env python

# ToDo:
#
# 1) Handle zip attachments

import os
import string
import sys
import errno
import mimetypes
import tempfile
import shutil
import re
import subprocess

from email.parser import Parser

from optparse import OptionParser

CUCKOO_SUBMIT = '/foo/bar'    # Path to submit.py
TMP_DIR = '/Users/kovar/Tools/cuckoo/tmp' # Path to create temp directory
DEBUG = True  # If True, keep the temporary directory after submitting

def submit_file(file, package):
    
    file = "\'" + file + "\'"
#    print CUCKOO_SUBMIT + ' ' + file + ' ' + package
    return_code = subprocess.call([CUCKOO_SUBMIT, file, package], shell=False)
    
def url_to_file(tmp_dir, url):
    
    try:
        (fd, name) = tempfile.mkstemp(dir=tmp_dir, text=True)
    except:
        print "Unable to create temporary file."
        sys.exit()
    
    try:
        f = os.fdopen(fd, "w+b")
    except:
        print "Unable to create temporary file."
        sys.exit()

    f.write('[InternetShortcut]\n')
    f.write('URL=' + url + '\n')
    f.close()
    
    submit_file(name, 'ie')
    
def unpack_message(tmp_dir, msg):
    
    counter = 1
    for part in msg.walk():
        urls = re.findall("(?P<url>https?://[^\s]+)", part.as_string())
        for url in urls:
            url_to_file(tmp_dir, url)
#            if url != None:
#                print url.group("url")
        
        # multipart/* are just containers
        if part.get_content_maintype() == 'multipart':
            continue
        # Applications should really sanitize the given filename so that an
        # email message can't be used to overwrite important files
        filename = part.get_filename()

        if not filename:
            ext = mimetypes.guess_extension(part.get_content_type())
            if not ext:
                # Use a generic bag-of-bits extension
                ext = '.bin'
            filename = 'part-%03d%s' % (counter, ext)
        else:
            temp, ext = os.path.splitext(filename)
            ext = string.lower(ext)
            
        counter += 1

        # Write the attachment out to the file
        full_path = os.path.join(tmp_dir, filename)
        fp = open(full_path, 'wb')
        fp.write(part.get_payload(decode=True))
        fp.close()
        
        if (ext == '.docx' or ext == '.doc'):
            submit_file(full_path, 'doc')
        elif (ext == '.xlsx' or ext == '.xls'):
            submit_file(full_path, 'xls')
        elif (ext == '.exe'):
            submit_file(full_path, 'exe')
        elif (ext == '.dll'):
            submit_file(full_path, 'dll')
        elif (ext == '.pdf'):
            submit_file(full_path, 'pdf')    
        
def main():
    
    parser = OptionParser(usage="""\
Unpack a MIME message into a directory of files.

Usage: %prog [options]
""")
    parser.add_option('-m', '--message',
                      type='string', action='store',
                      help="""Specify the MIME message to unpack.
                      """)
    opts, args = parser.parse_args()
    if not opts.message:
        parser.print_help()
        sys.exit(1)

    try:
        tmp_dir = tempfile.mkdtemp(dir=TMP_DIR)
    except:
        print "Unable to open temporary directory."
        sys.exit()
        
        
    # Open and parse the message to unpack        
    fp = open(opts.message)
    msg = Parser().parse(fp)
    fp.close()
    
    unpack_message(tmp_dir, msg)

    if DEBUG == False:  # Keep the tmp directory and its contents around if debugging
        try:
            shutil.rmtree(tmp_dir) # delete directory
        except OSError, e:
            if e.errno != 2: # code 2 - no such file or directory
                raise        
    
if __name__ == '__main__':
    main()
    
